package DyingBaby.backend.plan;

import DyingBaby.backend.catalog.Catalog;
import DyingBaby.backend.catalog.ColumnDef;
import DyingBaby.backend.catalog.ColumnType;
import DyingBaby.backend.catalog.TableSchema;
import DyingBaby.backend.sql.Ast;

import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

/**主要功能:
 * 1.接收一个AST语句（Ast.Stmt），根据其类型分发给具体的计划方法。
 * 2.对每种类型的语句进行语义检查（例如，表是否存在、列是否存在、类型是否匹配等）。
 * 3.生成对应的逻辑计划（LogicalPlan），这些逻辑计划是后续执行的基础。
 */
public class Planner {
    private final Catalog catalog;

    public Planner(Catalog catalog) { this.catalog = catalog; }

    //计划器的主要入口方法，根据语句类型分发到具体的计划方法
    public LogicalPlan plan(Ast.Stmt stmt) {
        if (stmt instanceof Ast.CreateTable) return planCreate((Ast.CreateTable) stmt);
        if (stmt instanceof Ast.Insert) return planInsert((Ast.Insert) stmt);
        if (stmt instanceof Ast.Select) return optimize(planSelect((Ast.Select) stmt));
        if (stmt instanceof Ast.Update) return planUpdate((Ast.Update) stmt);
        if (stmt instanceof Ast.Delete) return planDelete((Ast.Delete) stmt);
        throw new IllegalArgumentException("Unsupported statement for planning");
    }

    //语义检查：表不能已存在
    //语义检查：列名必须唯一
    //语义检查：主键列必须在列定义中
    //创建并返回 CreateTablePlan
    private LogicalPlan planCreate(Ast.CreateTable ct) {
        // semantic: table not exists; columns non-empty; names unique; types valid; pk in columns
        if (catalog.listTables().containsKey(ct.table)) {
            throw new IllegalArgumentException("Table exists: " + ct.table);
        }
        ensureColumnNamesUnique(ct.columns);
        if (ct.primaryKey != null) ensurePkInColumns(ct.primaryKey, ct.columns);
        return new CreateTablePlan(ct.table, ct.columns, ct.primaryKey);
    }

    //获取表结构信息
    // 语义检查：插入的值必须包含所有列（不支持部分列插入）
    //语义检查：值类型必须与列定义匹配
    //创建并返回 InsertPlan
    private LogicalPlan planInsert(Ast.Insert ins) {
        TableSchema schema = catalog.getTable(ins.table);
        // semantic: all columns present, types/length match, no extras
        Map<String, ColumnDef> cols = schema.getColumnMap();
        if (ins.values.size() != cols.size()) throw new IllegalArgumentException("Unexpected columns in record");
        for (ColumnDef cd : cols.values()) {
            if (!ins.values.containsKey(cd.getName())) throw new IllegalArgumentException("Missing column: " + cd.getName());
            checkValueType(cd, ins.values.get(cd.getName()));
        }
        return new InsertPlan(ins.table, new LinkedHashMap<>(ins.values));
    }

    //获取表结构信息
    //创建基础计划：顺序扫描 (SeqScanPlan)
    //如果有条件，添加过滤计划 (FilterPlan)
    //添加投影计划
    //返回完整的逻辑计划树
    private LogicalPlan planSelect(Ast.Select sel) {
        TableSchema schema = catalog.getTable(sel.table);
        LogicalPlan root = new SeqScanPlan(sel.table);
        if (sel.predicate != null) {
            validateExpression(sel.predicate, schema);
            root = new FilterPlan(root, sel.predicate);
        }
        // project all columns for now
        root = new ProjectPlan(root);
        return root;
    }

    // 加速索引：
    // 规则1：若是 Project( Filter( SeqScan(table) ) ) 且谓词为主键等值 col=pk，则重写为 PkLookup(table, value)
    private LogicalPlan optimize(LogicalPlan plan) {
        if (plan instanceof ProjectPlan) {
            LogicalPlan c1 = ((ProjectPlan) plan).child;
            if (c1 instanceof FilterPlan) {
                FilterPlan fp = (FilterPlan) c1;
                if (fp.child instanceof SeqScanPlan) {
                    SeqScanPlan sp = (SeqScanPlan) fp.child;
                    TableSchema schema = catalog.getTable(sp.table);
                    //优化条件(触发条件): 主键列必须存在
                    String pk = schema.getPrimaryKey();
                    //过滤条件变得复杂
                    if (pk != null && fp.predicate instanceof Ast.Eq) {
                        Ast.Eq eq = (Ast.Eq) fp.predicate;
                        if (pk.equals(eq.col)) {
                            return new PkLookupPlan(sp.table, eq.value);
                        }                    }
                }
            }
        }
        return plan;
    }

    //获取表结构信息
    //语义检查：SET 子句中的列必须存在
    //语义检查：SET 值类型必须匹配列定义
    //处理 WHERE 条件（目前支持复杂查询）
    //创建并返回 UpdatePlan
    private LogicalPlan planUpdate(Ast.Update up) {
        TableSchema schema = catalog.getTable(up.table);
        // check SET columns exist and value type
        for (Map.Entry<String,Object> e : up.set.entrySet()) {
            ColumnDef cd = schema.getColumnMap().get(e.getKey());
            if (cd == null) throw new IllegalArgumentException("Unknown column in SET: " + e.getKey());
            checkValueType(cd, e.getValue());
        }
        if (up.predicate != null) {
            validateExpression(up.predicate, schema);
        }

        return new UpdatePlan(up.table, new LinkedHashMap<>(up.set), up.predicate);
    }

    //获取表结构信息
    //处理 WHERE 条件（支持复杂条件)
    //创建并返回 DeletePlan
    private LogicalPlan planDelete(Ast.Delete del) {
        TableSchema schema = catalog.getTable(del.table);
        if (del.predicate != null) {
            validateExpression(del.predicate, schema);
        }
        return new DeletePlan(del.table, del.predicate);
    }

    //以下全是辅助方法
    //确保列名唯一
    private void ensureColumnNamesUnique(List<ColumnDef> cols) {
        java.util.Set<String> s = new java.util.HashSet<>();
        for (ColumnDef cd : cols) {
            if (!s.add(cd.getName())) throw new IllegalArgumentException("Duplicate column: " + cd.getName());
        }
    }


    //确保主键列存在
    private void ensurePkInColumns(String pk, List<ColumnDef> cols) {
        for (ColumnDef cd : cols) if (cd.getName().equals(pk)) return;
        throw new IllegalArgumentException("PRIMARY KEY not found in columns: " + pk);
    }

    // 检查值类型是否与列定义匹配
    private void checkValueType(ColumnDef cd, Object v) {
        if (cd.getType() == ColumnType.INT) {
            if (!(v instanceof Number)) throw new IllegalArgumentException("Column " + cd.getName() + " expects INT");
        } else if (cd.getType() == ColumnType.STRING) {
            if (!(v instanceof String)) throw new IllegalArgumentException("Column " + cd.getName() + " expects STRING");
            Integer ml = cd.getMaxLength();
            if (ml != null && ((String) v).length() > ml) throw new IllegalArgumentException("Column " + cd.getName() + " exceeds max length " + ml);
        }
    }

    // 验证表达式的语义正确性
    private void validateExpression(Ast.Expr expr, TableSchema schema) {
        if (expr instanceof Ast.Eq) {
            Ast.Eq eq = (Ast.Eq) expr;
            validateColumnExists(eq.col, schema);
        } else if (expr instanceof Ast.Ne) {
            Ast.Ne ne = (Ast.Ne) expr;
            validateColumnExists(ne.col, schema);
        } else if (expr instanceof Ast.Gt) {
            Ast.Gt gt = (Ast.Gt) expr;
            validateColumnExists(gt.col, schema);
        } else if (expr instanceof Ast.Lt) {
            Ast.Lt lt = (Ast.Lt) expr;
            validateColumnExists(lt.col, schema);
        } else if (expr instanceof Ast.Ge) {
            Ast.Ge ge = (Ast.Ge) expr;
            validateColumnExists(ge.col, schema);
        } else if (expr instanceof Ast.Le) {
            Ast.Le le = (Ast.Le) expr;
            validateColumnExists(le.col, schema);
        } else if (expr instanceof Ast.And) {
            Ast.And and = (Ast.And) expr;
            validateExpression(and.left, schema);
            validateExpression(and.right, schema);
        } else if (expr instanceof Ast.Or) {
            Ast.Or or = (Ast.Or) expr;
            validateExpression(or.left, schema);
            validateExpression(or.right, schema);
        } else if (expr instanceof Ast.Not) {
            Ast.Not not = (Ast.Not) expr;
            validateExpression(not.expr, schema);
        } else if (expr instanceof Ast.Paren) {
            Ast.Paren paren = (Ast.Paren) expr;
            validateExpression(paren.expr, schema);
        } else {
            throw new IllegalArgumentException("Unsupported expression type: " + expr.getClass().getSimpleName());
        }
    }

    // 验证列是否存在
    private void validateColumnExists(String col, TableSchema schema) {
        if (!schema.getColumnMap().containsKey(col)) {
            throw new IllegalArgumentException("Unknown column: " + col);
        }
    }
} 